#!/bin/bash
set -euo pipefail

# 脚本名称：dir2qcow2.sh

# 功能说明：
# 本脚本用于将目录转换为 QCOW2 格式的虚拟磁盘镜像。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-nbd、mkfs.ext4、numfmt、du、awk、parted 等）。
# 2. 计算目录大小并创建虚拟磁盘。
# 3. 选择可用的 NBD 设备，并将虚拟磁盘挂载到 NBD 设备。
# 4. 创建分区并格式化为 ext4 文件系统。
# 5. 将目录内容复制到虚拟磁盘。
# 6. 清理挂载的 NBD 设备。

# 注意事项：
# - 需要 root 权限运行部分命令，如 qemu-nbd、mkfs.ext4、parted。
# - 在执行操作前，建议备份重要数据，以防万一。

#--------- 依赖管理 ---------#
install_dependencies() {
    local required=(qemu-nbd mkfs.ext4 numfmt du awk parted)
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -gt 0 ]; then
        echo "缺少依赖: ${missing[*]}，正在安装..."
        sudo apt-get update -qq
        sudo apt-get install -y -qq qemu-utils e2fsprogs coreutils gawk parted >/dev/null
    fi
}

#--------- 全局变量 ---------#
nbd_device=""
virtual_disk_attached=false
readonly SAFETY_SIZE="1G"
readonly MOUNT_POINT="/mnt/temp_disk"

# 清理函数：断开当前连接的 NBD 设备
cleanup() {
    echo ">>> 执行清理操作..."
    if mountpoint -q "$MOUNT_POINT"; then
        sudo umount "$MOUNT_POINT" || true
    fi
    if [[ -d "$MOUNT_POINT" ]]; then
        sudo rmdir "$MOUNT_POINT" || true
    fi
    if lsblk "$nbd_device" &>/dev/null; then
        sudo qemu-nbd -d "$nbd_device" || true
    fi
    virtual_disk_attached=false
}
trap cleanup EXIT SIGINT SIGTERM

#--------- NBD设备管理 ---------#
# 从 /dev/nbd* 中选择一个空闲的设备
select_nbd() {
    for dev in /dev/nbd*; do
        # 检查设备是否存在对应的 pid 文件且内容为空（表示未被占用）
        pid_file="/sys/block/$(basename "$dev")/pid"
        if [ -f "$pid_file" ]; then
            if [ ! -s "$pid_file" ]; then
                echo "$dev"
                return 0
            fi
        else
            # 若 pid 文件不存在，则认为设备可用
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 挂载虚拟磁盘到NBD设备
attach_disk() {
    local disk_image="$1"
    sudo modprobe nbd max_part=8
    nbd_device=$(select_nbd)
    if [ -z "$nbd_device" ]; then
        echo "无可用NBD设备" >&2
        exit 1
    fi

    echo "正在附加 $disk_image 到 $nbd_device..."
    sudo qemu-nbd -c "$nbd_device" "$disk_image" || { echo "附加失败"; exit 1; }
    sudo udevadm settle
    virtual_disk_attached=true
}

#--------- 磁盘操作 ---------#
# 计算目录大小并添加安全余量
calculate_disk_size() {
    local dir="$1"
    local dir_size=$(sudo du -sb "$dir" | awk '{print $1}')
    local safety_bytes=$(numfmt --from=iec "$SAFETY_SIZE")
    echo $((dir_size + safety_bytes))
}

# 创建虚拟磁盘
create_virtual_disk() {
    local disk_path="$1" size="$2"
    echo ">>> 创建虚拟磁盘: $disk_path (大小: $(numfmt --to=iec "$size"))"
    qemu-img create -f qcow2 "$disk_path" "$size"
}

# 准备块设备
prepare_block_device() {
    echo ">>> 创建分区表"
    sudo parted "$nbd_device" mklabel msdos
    sudo parted -a opt "$nbd_device" mkpart primary ext4 0% 100%
    
    local partition="${nbd_device}p1"
    echo ">>> 格式化分区: $partition"
    sudo mkfs.ext4 "$partition"
}

# 复制数据
copy_data() {
    local src="$1"
    echo ">>> 挂载分区到: $MOUNT_POINT"
    sudo mkdir -p "$MOUNT_POINT"
    sudo mount "${nbd_device}p1" "$MOUNT_POINT"
    
    echo ">>> 复制数据: $src => $MOUNT_POINT"
    sudo cp -a "$src/." "$MOUNT_POINT/"
}

#--------- 主流程 ---------#
main() {
    install_dependencies

    echo "=== 目录转虚拟磁盘工具 ==="
    
    # 参数验证
    if [ $# -ne 2 ]; then
        echo "用法: $0 <源目录> <输出镜像路径>"
        echo "示例: $0 /home/user/data ./data.qcow2"
        exit 1
    fi
    
    local src_dir="$1"
    local disk_path="$2"
    
    # 大小计算
    local disk_size=$(calculate_disk_size "$src_dir")
    
    # 创建磁盘
    create_virtual_disk "$disk_path" "$disk_size"
    
    # 处理块设备
    attach_disk "$disk_path"
    prepare_block_device
    
    # 数据复制
    copy_data "$src_dir"
    
    echo "=== 转换完成 ==="
    echo "输出镜像: $disk_path"
    echo "可使用以下命令挂载检查:"
    echo "  qemu-nbd -c $nbd_device $disk_path && mount ${nbd_device}p1 /mnt"
}

main "$@"